package com.castle.fortress.admin.ai.config;

import com.castle.fortress.admin.ai.llm.Llm;
import com.castle.fortress.admin.ai.message.AiMessage;
import com.castle.fortress.admin.ai.message.HumanMessage;
import com.castle.fortress.admin.ai.prompt.HistoriesPrompt;

import java.util.Scanner;

public class QwenTest {

    public static void main(String[] args) throws InterruptedException {
        QwenLlmConfig config = new QwenLlmConfig();
        config.setApiKey("sk-887ef1c8383e4027b014f85372107fda");
        config.setModel("qwen-max");

        Llm llm = new QwenLlm(config);
        HistoriesPrompt prompt = new HistoriesPrompt();

        System.out.println("您想问什么？");
        Scanner scanner = new Scanner(System.in);
        String userInput = scanner.nextLine();

        while (userInput != null) {

            prompt.addMessage(new HumanMessage(userInput));

            llm.chatAsync(prompt, (context, response) -> {
                AiMessage message = response.getMessage();
                System.out.println(message.getStatus());
                System.out.println(">>>> " + message.getContent());
            });

            userInput = scanner.nextLine();
        }

//
//        llm.chatAsync("请写一个小兔子战胜大灰狼的故事", (context, response) -> {
//            AiMessage message = response.getMessage();
//            System.out.println(">>>> " + message.getContent() + "   （消耗的 Token 量：" + message.getTotalTokens() + "）");
//        });
//
//        Thread.sleep(10000);
    }
}
